#!/home/agilex/miniconda3/envs/graspnet/bin/python
# -- coding: UTF-8
import rospy
import std_msgs.msg
from tf2_ros import TransformBroadcaster
from geometry_msgs.msg import TransformStamped
import numpy as np


import struct
import std_msgs
from sensor_msgs.msg import PointCloud2, PointField
import sensor_msgs.point_cloud2 as pc2

import pyrealsense2 as rs

import queue
import threading
import open3d as o3d

intrinsics_matrix = np.array([[609.0, 0., 320.],
                         [ 0.,609.0, 240],
                         [0,0,1]], dtype=np.float32)

factor_depth = np.array([[1000.]])

class CameraInfo():
    def __init__(self, width, height, fx, fy, cx, cy, scale):
        self.width = width
        self.height = height
        self.fx = fx
        self.fy = fy
        self.cx = cx
        self.cy = cy
        self.scale = scale

def rad2deg(rad):
    return rad / np.pi * 180. 


def rgbd2cloud(color, depth, camera, rgb_flag = False, num_point = 10000):
    assert(depth.shape[0] == camera.height and depth.shape[1] == camera.width)
    xmap = np.arange(camera.width)
    ymap = np.arange(camera.height)
    xmap, ymap = np.meshgrid(xmap, ymap)  # 创建图像网格
    points_z = depth / camera.scale       # 深度信息  使用了广播机制
   
    points_x = (xmap - camera.cx) * points_z / camera.fx  # X = (u - cx) * Z /fx
    points_y = (ymap - camera.cy) * points_z / camera.fy  # Y = (v - cx) * Y /fy
    cloud = np.stack([points_x, points_y, points_z], axis=-1)  # 沿着最后一个轴增加维度
    if rgb_flag:
        color = color[:, :, ::-1]   # bgr2rgb
        cloud = np.concatenate((cloud, color), axis=2)
        # print("cloud: ", cloud.shape)
        cloud = cloud.reshape([-1, 6])
    else:
        cloud = cloud.reshape([-1, 3])
    cloud = cloud[(cloud[:, 2] > 0.2) & (cloud[:, 2] < 2)]
    
    cloud_size = cloud.shape[0]
    if(cloud_size >= num_point):
        idxs = np.random.choice(cloud.shape[0] , num_point, replace=False)
    else:
        idxs1 = np.arange(cloud_size)  # 补点
        idxs2 = np.random.choice(cloud_size, num_point - cloud_size, replace=True)
        idxs = np.concatenate([idxs1, idxs2], axis=0)
        idxs = np.random.choice(cloud.shape[0] , num_point, replace=False)
    cloud = cloud[idxs]
    
    return cloud

def float_to_rgb(float_rgb):
    """
    Convert a float RGB to separate r, g, b values.
    
    Args:
    - float_rgb: A float representing RGB values packed into a single float.
    
    Returns:
    - A tuple (r, g, b) where r, g, b are integers in the range [0, 255].
    """
    # Convert float to 32-bit packed integer
    s = struct.pack('>f', float_rgb)
    int_rgb = struct.unpack('>l', s)[0]
    # Extract RGB components
    r = (int_rgb >> 16) & 0xFF
    g = (int_rgb >> 8) & 0xFF
    b = int_rgb & 0xFF
    return r, g, b

def pointcloud2_to_numpy(pointcloud_msg, num_points=10000):
    """
    Convert a ROS PointCloud2 message to a numpy array.
    
    Args:
    - pointcloud_msg: sensor_msgs/PointCloud2 message
    
    Returns:
    - numpy array of shape (n, 6) with each row as (x, y, z, r, g, b)
    """
    # Extract points from point cloud message
    point_list = []
    for point in pc2.read_points(pointcloud_msg, skip_nans=True, field_names=("x", "y", "z", "rgb")):
        x, y, z, float_rgb = point
        r, g, b = float_to_rgb(float_rgb)
        point_list.append([x, y, z, r, g, b])
    
    # Convert to numpy array
    cloud = np.array(point_list, dtype=np.float32)
    
    cloud_size = cloud.shape[0]
    if(cloud_size >= num_points):
        idxs = np.random.choice(cloud_size , num_points, replace=False)
    else:
        idxs1 = np.arange(cloud_size)  # 补点
        idxs2 = np.random.choice(cloud_size, num_points - cloud_size, replace=True)
        idxs = np.concatenate([idxs1, idxs2], axis=0)
        idxs = np.random.choice(cloud_size , num_points, replace=False)
    cloud = cloud[idxs]

    return cloud


def open3d2msg(open3d_cloud, frame_id="map"):
    points = np.asarray(open3d_cloud.points, dtype=np.float32)
    colors = (np.asarray(open3d_cloud.colors) * 255).astype(np.uint8) if len(open3d_cloud.colors) > 0 else np.full((points.shape[0], 3), 255, dtype=np.uint8)

    # Prepare data for PointCloud2
    cloud_data = np.zeros((points.shape[0], 4), dtype=np.float32)
    cloud_data[:, 0:3] = points

    for i in range(colors.shape[0]):
        r, g, b = colors[i]
        rgb_packed = (r << 16) | (g << 8) | b
        rgb_float = struct.unpack('f', struct.pack('I', rgb_packed))[0]
        cloud_data[i, 3] = rgb_float

    # Flatten data to byte array
    cloud_data = cloud_data.flatten().tobytes()

    # Create PointCloud2 message
    header = std_msgs.msg.Header()
    header.stamp = rospy.Time.now()
    header.frame_id = frame_id

    fields = [
        PointField('x', 0, PointField.FLOAT32, 1),
        PointField('y', 4, PointField.FLOAT32, 1),
        PointField('z', 8, PointField.FLOAT32, 1),
        PointField('rgb', 12, PointField.FLOAT32, 1),
    ]

    point_cloud_msg = PointCloud2(
        header=header,
        height=1,
        width=points.shape[0],
        is_dense=False,
        is_bigendian=False,
        fields=fields,
        point_step=16,
        row_step=16 * points.shape[0],
        data=cloud_data
    )
    return point_cloud_msg


def convert_numpy_to_pointcloud2(np_array):
    """
    Convert numpy array to sensor_msgs/PointCloud2.
    
    np_array: A numpy array of shape (n, 6) with each row as (x, y, z, r, g, b)
    
    Returns:
    PointCloud2 message
    """
    header = std_msgs.msg.Header()
    header.stamp = rospy.Time.now()
    header.frame_id = "camera_color_frame"

    fields = [
        PointField('x', 0, PointField.FLOAT32, 1),
        PointField('y', 4, PointField.FLOAT32, 1),
        PointField('z', 8, PointField.FLOAT32, 1),
        PointField('rgba', 12, PointField.UINT32, 1),
    ]

    point_cloud = []
    for point in np_array:
        x, y, z, r, g, b = point
        r = min(max(int(r), 0), 255)
        g = min(max(int(g), 0), 255)
        b = min(max(int(b), 0), 255)
        a = 255  # Full opacity
        rgba = struct.unpack('I', struct.pack('BBBB', b, g, r, a))[0]
        point_cloud.append([x, y, z, rgba])

    pc2_msg = pc2.create_cloud(header, fields, point_cloud)
    
    return pc2_msg


class PublishNode:
    def __init__(self) -> None:
        rospy.init_node("grasp_node")
        self.tf_broad = TransformBroadcaster()
        self.msg_tf = TransformStamped()
        self.rate = rospy.Rate(30)
        self.pub_cloud = rospy.Publisher("/points", PointCloud2, queue_size=10)
        self.camera = CameraInfo(640, 480, intrinsics_matrix[0][0], intrinsics_matrix[1][1], intrinsics_matrix[0][2], intrinsics_matrix[1][2], factor_depth)
        
        self.cloudQue = queue.Queue()
        self.cloud = None
        self.num_points = 10000

        t1 = threading.Thread(target=self.process_cloud)
        t1.start()
        
        self.sub_cloud = rospy.Subscriber("/camera/depth/color/points", PointCloud2, self.callback, queue_size=10)

        rospy.spin()

    def callback(self, msg):
        self.cloudQue.put(msg)
        if(self.cloudQue.qsize()) > 10:
            self.cloudQue.get()
        print(self.cloudQue.qsize())


    def process_cloud(self):
        while not rospy.is_shutdown():
            if self.cloudQue.empty():
                continue
            cloud_msg = self.cloudQue.get()
            points = pointcloud2_to_numpy(cloud_msg, num_points=10000)
            
            # self.pub_cloud.publish(convert_numpy_to_pointcloud2(points))
            self.rate.sleep()


    def publish_tf(self, frame_id, child_frame_id, q, t):
        self.msg_tf.header.frame_id = frame_id
        self.msg_tf.child_frame_id = child_frame_id
        # -0.002 -0.015  0.201 -0.295 -0.293 -0.645  0.641
        self.msg_tf.transform.rotation.x = q[0]
        self.msg_tf.transform.rotation.y = q[1]
        self.msg_tf.transform.rotation.z = q[2]
        self.msg_tf.transform.rotation.w = q[3]
        self.msg_tf.transform.translation.x = t[0]
        self.msg_tf.transform.translation.y = t[1]
        self.msg_tf.transform.translation.z = t[2]
       
        self.msg_tf.header.stamp = rospy.Time.now()
        self.tf_broad.sendTransform(self.msg_tf)
            
    def publish_cloud(self, cloud):
        while not rospy.is_shutdown():
            # print(cloud)
            # msg_cloud = open3d2msg(cloud)
            msg_cloud = convert_numpy_to_pointcloud2(cloud)
            self.pub_cloud.publish(msg_cloud)
            self.rate.sleep()

if __name__ == "__main__":
    pub_node = PublishNode()